package com.glodon.paas.registry.manager;

import com.glodon.paas.registry.util.X509CertUtils;
import com.google.common.io.Files;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.math.BigInteger;
import java.net.URL;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

import static com.google.common.base.Strings.isNullOrEmpty;

/**
 * X.509 certificate manager
 *
 * @author Don Li
 */
public class X509CertManager {

    static {
        if (Security.getProvider(BouncyCastleProvider.PROVIDER_NAME) == null) {
            Security.addProvider(new BouncyCastleProvider());
        }
    }

    private static Logger LOGGER = LoggerFactory.getLogger(X509CertManager.class);

    private File keyFile;
    private String keyFilePwd;
    private String keyPwd;
    private long certTtl;
    private int keyLength;
    private String keyAlgId;
    private String keySeed;
    private String[] aliases;

    private AtomicInteger serialNum = new AtomicInteger(100);
    private Map<String, X509Certificate> certs = new HashMap<String, X509Certificate>();
    private Map<String, PrivateKey> keys = new HashMap<String, PrivateKey>();

    private File workingDir = Files.createTempDir();

    public X509CertManager(String keyFilePath, String keyFilePwd, String[] aliases, String keyPwd, long certTtl, int keyLength, String keyAlgId, String keySeed) throws Exception {
        if (keyFilePath.startsWith("classpath:")) {
            URL url = X509CertManager.class.getClassLoader().getResource(keyFilePath.substring(10));
            if (url != null) {
                keyFile = new File(url.toURI());
            } else {
                throw new FileNotFoundException("keystore is not found in classpath");
            }
        } else {
            keyFile = new File(keyFilePath);
        }
        this.keyFilePwd = keyFilePwd;
        this.aliases = aliases;
        this.keyPwd = keyPwd;
        this.certTtl = certTtl;
        this.keyLength = keyLength;
        this.keyAlgId = keyAlgId;
        this.keySeed = keySeed;
    }

    public void init() {
        try {
            if (!keyFile.exists()) {
                initCerts();
            }
            KeyStore keyStore = X509CertUtils.loadKeyStore(keyFile, keyFilePwd.toCharArray());
            for (String alias : aliases) {
                X509Certificate cert = (X509Certificate) keyStore.getCertificate(alias);
                certs.put(alias, cert);
                PrivateKey key = (PrivateKey) keyStore.getKey(alias, keyPwd.toCharArray());
                keys.put(alias, key);
            }
        } catch (Exception e) {
            LOGGER.error("Fail to load certificates and keys from keystore", e);
        }
        LOGGER.debug("X509CertManager has been initialized");
    }

    public File createKeystoreForInstanceCertChain(String serviceId, String instanceId, String keyPwd, String filePwd) {
        if (isNullOrEmpty(keyPwd) || isNullOrEmpty(filePwd) || isNullOrEmpty(instanceId)) {
            return null;
        }
        try {
            File resultFile = File.createTempFile("cert-", ".pkcs", workingDir);
            X509Certificate caCert = getMasterCert();
            X509Certificate serviceCert = getServiceCert(serviceId);

            if (caCert == null || serviceCert == null) {
                return null;
            }

            Date notBefore = new Date(System.currentTimeMillis());
            Date notAfter = new Date(System.currentTimeMillis() + certTtl);

            KeyPairGenerator keyGen = KeyPairGenerator.getInstance(keyAlgId);
            keyGen.initialize(keyLength, new SecureRandom(keySeed.getBytes()));
            KeyPair keyPair = keyGen.generateKeyPair();
            PrivateKey privateKey = keyPair.getPrivate();
            PublicKey publicKey = keyPair.getPublic();

            String subject = "C=CN,O=Glodon,OU=PASS,CN=" + instanceId;

            X509Certificate cliCert = X509CertUtils.createCert(serviceCert, subject, notBefore, notAfter, publicKey, getPrivateKey(serviceId), serviceCert.getPublicKey(), getNextSerialNumber());

            X509Certificate[] chain = new X509Certificate[3];
            chain[2] = caCert;
            chain[1] = serviceCert;
            chain[0] = cliCert;

            X509CertUtils.storeChainAsPkcs12(chain, privateKey, keyPwd.toCharArray(), instanceId, resultFile, filePwd.toCharArray());
            return resultFile;
        } catch (Exception e) {
            LOGGER.error("Fail to create cert for instance", e);
        }
        return null;
    }

    public File getOrCreateKeystoreForServiceCert(String serviceId, String filePwd) {
        try {
            File resultFile = File.createTempFile("cert-", ".pkcs", workingDir);
            X509Certificate serviceCert = getServiceCert(serviceId);

            if (serviceCert == null) {
                return null;
            }

            X509CertUtils.storeCertAsPkcs12(serviceCert, serviceId, resultFile, filePwd.toCharArray());
            return resultFile;
        } catch (Exception e) {
            LOGGER.error("Fail to retrieve cert for service", e);
        }
        return null;
    }

    private void initCerts() {
        try {
            KeyStore store = KeyStore.getInstance("PKCS12", BouncyCastleProvider.PROVIDER_NAME);
            store.load(null, null);

            KeyPairGenerator keyGen = KeyPairGenerator.getInstance(keyAlgId);
            keyGen.initialize(keyLength, new SecureRandom(keySeed.getBytes()));
            KeyPair keyPair = keyGen.generateKeyPair();
            PrivateKey caPrivKey = keyPair.getPrivate();
            PublicKey caPubKey = keyPair.getPublic();
            Date notBefore = new Date(System.currentTimeMillis());
            Date notAfter = new Date(System.currentTimeMillis() + certTtl * 3);

            // create master cert
            X509Certificate caCert = initMasterCert(store, caPrivKey, caPubKey, notBefore, notAfter);
            // create service cert for document and job
            addServiceCert(store, keyGen, "document", caCert, caPrivKey);
            addServiceCert(store, keyGen, "job", caCert, caPrivKey);

            // output to file
            store.store(new FileOutputStream(keyFile), keyFilePwd.toCharArray());
        } catch (Exception e) {
            LOGGER.error("Fail to build keystore", e);
        }
    }

    private X509Certificate initMasterCert(KeyStore store, PrivateKey caPrivKey, PublicKey caPubKey, Date notBefore, Date notAfter) throws Exception {
        String issuer = "C=CN,O=Glodon,OU=PASS";
        String subject = "C=CN,O=Glodon,OU=PASS";

        X509Certificate caCert = X509CertUtils.createMasterCert(issuer, caPrivKey, subject, caPubKey, notBefore, notAfter, BigInteger.valueOf(1));
        Certificate[] chain = new Certificate[1];
        chain[0] = caCert;
        store.setKeyEntry("master", caPrivKey, keyPwd.toCharArray(), chain);
        return caCert;
    }

    private void addServiceCert(KeyStore store, KeyPairGenerator keyGen, String serviceId, X509Certificate caCert, PrivateKey caPrivKey) throws Exception {
        KeyPair keyPair = keyGen.generateKeyPair();
        PrivateKey intPrivKey = keyPair.getPrivate();
        PublicKey intPubKey = keyPair.getPublic();
        String intSubject = "C=CN,O=Glodon,OU=PASS,CN=" + serviceId;
        Date notBefore = new Date(System.currentTimeMillis());
        Date notAfter = new Date(System.currentTimeMillis() + certTtl);
        X509Certificate intCert = X509CertUtils.createIntermediateCert(caCert, caPrivKey, intSubject, intPubKey, notBefore, notAfter, BigInteger.valueOf(10));
        Certificate[] chain = new Certificate[2];
        chain[1] = caCert;
        chain[0] = intCert;
        store.setKeyEntry(serviceId, intPrivKey, keyPwd.toCharArray(), chain);
    }

    private X509Certificate getMasterCert() {
        return certs.get("master");
    }

    private X509Certificate getServiceCert(String serviceId) {
        if (serviceId.equals("master")) {
            return null; // do not expose the master key
        }
        return certs.get(serviceId);
    }

    private PrivateKey getPrivateKey(String serviceId) {
        return keys.get(serviceId);
    }

    private BigInteger getNextSerialNumber() {
        return BigInteger.valueOf(serialNum.getAndIncrement());
    }
}
